{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Long-Term Memory\n",
    "\n",
    "![Short-term Memory Summarization](images/short-term%20memory%20summarization.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Keeping long-term memories about collaboration \n",
    "\n",
    "![Short-term Memory vs. Long-term Memory](images/short-term%20memory%20vs%20long-term%20memory.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LangGraph Store\n",
    "\n",
    "- Namespaces (User Profile, User Memories, User Preferences)\n",
    "    - key-value pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.store.memory import InMemoryStore\n",
    "\n",
    "namespace = (\"evgeny\", \"facts\")\n",
    "\n",
    "# Save a memory to namespace as key and value\n",
    "key = \"fun_fact_1\"\n",
    "\n",
    "# The value needs to be a dictionary  \n",
    "value = {\n",
    "  \"fact\": \"Evgeny likes to ask philosophical questions late at night.\",\n",
    "  \"source\": \"conversation_2024_05_30\"\n",
    "}\n",
    "\n",
    "# Save the memory\n",
    "in_memory_store = InMemoryStore()\n",
    "in_memory_store.put(namespace, key, value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### retreive all stored memos for namespace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "memories = in_memory_store.search(namespace)\n",
    "memories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Metatdata \n",
    "memories[0].dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find by namespace and key\n",
    "memory = in_memory_store.get(namespace, key)\n",
    "memory.dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chat-bot with short and long memory\n",
    "\n",
    "### Short-term memory - keeps info about current conversation, uses Checkpoints\n",
    "\n",
    "### Long-term memory - keeps info about user's profile, preferences, etc., uses MemoryStore\n",
    "\n",
    "#### Configuration\n",
    "```\n",
    "config = {\n",
    "    \"configurable\": {\n",
    "        \"thread_id\": \"1\", \n",
    "        \"user_id\": \"1\"\n",
    "    }\n",
    "}\n",
    "```\n",
    "\n",
    "#### Store\n",
    "```\n",
    "namespace = (\"memory\", user_id)\n",
    "key = \"user_details\"\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, display\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import StateGraph, MessagesState, START, END\n",
    "from langchain_core.runnables.config import RunnableConfig\n",
    "from langgraph.store.base import BaseStore\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "### Nodes\n",
    "\n",
    "def chat(state: MessagesState, config: RunnableConfig, store: BaseStore):\n",
    "    user_id = config[\"configurable\"][\"user_id\"]\n",
    "\n",
    "    # Retrieve memory from the store\n",
    "    user_details = store.get((\"memory\", user_id), \"user_details\")\n",
    "\n",
    "    # Extract the actual memory content if it exists and add a prefix\n",
    "    if user_details:\n",
    "        # Value is a dictionary with a memory key\n",
    "        user_details_content = user_details.value.get('memory')\n",
    "    else:\n",
    "        user_details_content = \"No existing details found.\"\n",
    "\n",
    "    # Format the memory in the system prompt\n",
    "    system_msg = f\"\"\"\n",
    "    You are a helpful assistant with memory capabilities.\n",
    "    If user-specific memory is available, use it to personalize \n",
    "    your responses based on what you know about the user.\n",
    "    \n",
    "    Your goal is to provide relevant, friendly, and tailored \n",
    "    assistance that reflects the user’s preferences, context, and past interactions.\n",
    "\n",
    "    If the user’s name or relevant personal context is available, always personalize your responses by:\n",
    "        – Addressing the user by name (e.g., \"Sure, Bob...\") when appropriate\n",
    "        – Referencing known projects, tools, or preferences (e.g., \"your MCP  server typescript based project\")\n",
    "        – Adjusting the tone to feel friendly, natural, and directly aimed at the user\n",
    "\n",
    "    Avoid generic phrasing when personalization is possible. For example, instead of \"In TypeScript apps...\" say \"Since your project is built with TypeScript...\"\n",
    "\n",
    "    Use personalization especially in:\n",
    "        – Greetings and transitions\n",
    "        – Help or guidance tailored to tools and frameworks the user uses\n",
    "        – Follow-up messages that continue from past context\n",
    "\n",
    "    Always ensure that personalization is based only on known user details and not assumed.\n",
    "    \n",
    "    The user’s memory (which may be empty) is provided as: {user_details_content}\n",
    "    \"\"\"\n",
    "    \n",
    "    response = model.invoke([SystemMessage(content=system_msg)] + state[\"messages\"])\n",
    "\n",
    "    return {\"messages\": response}\n",
    "\n",
    "\n",
    "def update_memory(state: MessagesState, config: RunnableConfig, store: BaseStore):\n",
    "    user_id = config[\"configurable\"][\"user_id\"]\n",
    "    namespace = (\"memory\", user_id)\n",
    "    key = \"user_details\"\n",
    "    user_details = store.get(namespace, key)\n",
    "        \n",
    "    if user_details:\n",
    "        user_details_content = user_details.value.get('memory')\n",
    "    else:\n",
    "        user_details_content = \"No existing details found.\"\n",
    "\n",
    "    # Format the memory in the system prompt\n",
    "    system_msg = f\"\"\"\n",
    "    You are responsible for updating and maintaining accurate user memory to enable personalized responses.\n",
    "\n",
    "    CURRENT USER DETAILS:\n",
    "    {user_details_content}\n",
    "\n",
    "    INSTRUCTIONS:\n",
    "    1. Carefully review the chat history below.\n",
    "    2. Identify any new, explicitly stated user information, such as:\n",
    "        - Personal details (e.g., name, location)\n",
    "        - Preferences (likes, dislikes)\n",
    "        - Interests and hobbies\n",
    "        - Experiences or background\n",
    "        - Goals and future plans\n",
    "    3. If no new information is present, do not output anything.\n",
    "    4. If new information is found:\n",
    "        - Merge it with the existing memory\n",
    "        - Format the updated memory as a clear, bulleted list\n",
    "        - Include only factual, user-stated details\n",
    "    5. If new information contradicts existing memory, keep the most recent version stated by the user.\n",
    "\n",
    "    Important:\n",
    "    - Do NOT include summaries like \"no update needed\".\n",
    "    - ONLY return output when actual new user information is added.\n",
    "\n",
    "    Your final output should either be a clean, updated bulleted list — or nothing at all.\n",
    "    \"\"\"\n",
    "    \n",
    "    new_memory = model.invoke([SystemMessage(content=system_msg)] + state['messages'])\n",
    "\n",
    "    if new_memory.content.strip():\n",
    "        store.put(namespace, key, {\"memory\": new_memory.content})\n",
    "\n",
    "\n",
    "# Define the graph\n",
    "builder = StateGraph(MessagesState)\n",
    "builder.add_node(\"chat\", chat)\n",
    "builder.add_node(\"update_memory\", update_memory)\n",
    "builder.add_edge(START, \"chat\")\n",
    "builder.add_edge(\"chat\", \"update_memory\")\n",
    "builder.add_edge(\"update_memory\", END)\n",
    "\n",
    "long_term_memory = InMemoryStore()\n",
    "short_term_memory = MemorySaver()\n",
    "\n",
    "graph = builder.compile(checkpointer=short_term_memory, store=long_term_memory)\n",
    "\n",
    "display(Image(graph.get_graph(xray=1).draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start a new conversation\n",
    "thread = {\"configurable\": {\"thread_id\": \"1\", \"user_id\": \"1\"}}\n",
    "\n",
    "# define intiial user request\n",
    "initial_input = {\"messages\": HumanMessage(content=\"\"\"\n",
    "Hey, how are you? I am Evgeny, I am a software engineer and I need your help with my project. \n",
    "This is a TaskManager Java Spring Boot application.\n",
    "\"\"\")}\n",
    "\n",
    "# run the graph and stream in values mode\n",
    "for event in graph.stream(initial_input, thread, stream_mode=\"values\"):\n",
    "    event['messages'][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_id = thread[\"configurable\"][\"user_id\"]\n",
    "namespace = (\"memory\", user_id)\n",
    "key = \"user_details\"\n",
    "memory = long_term_memory.get(namespace, key)\n",
    "memory.dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let's start another thread"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start a new conversation\n",
    "thread = {\"configurable\": {\"thread_id\": \"2\", \"user_id\": \"1\"}}\n",
    "\n",
    "# define intiial user request\n",
    "initial_input = {\"messages\": HumanMessage(content=\"\"\"\n",
    "How can I set up security configurations?\n",
    "\"\"\")}\n",
    "\n",
    "# run the graph and stream in values mode\n",
    "for event in graph.stream(initial_input, thread, stream_mode=\"values\"):\n",
    "    event['messages'][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_id = thread[\"configurable\"][\"user_id\"]\n",
    "namespace = (\"memory\", user_id)\n",
    "key = \"user_details\"\n",
    "memory = long_term_memory.get(namespace, key)\n",
    "memory.dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LangGraph Studio\n",
    "\n",
    "url: https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024\n",
    "\n",
    "prompt1: Hey, how are you? I am Evgeny, I am a software engineer and I need your help with my project. This is a TaskManager Java Spring Boot application.\n",
    "\n",
    "prompt2: How can I set up security configurations?\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
